import datetime
from django.core.management.base import BaseCommand
from django.db.models import Avg, Max, Min, Count, Sum, F
from tushare_data.models import StockDaily, StockBasic
from data_analysis.models import StockPerformance, SectorPerformance, MarketSummary


class Command(BaseCommand):
    help = '分析市场数据并生成统计信息'

    def add_arguments(self, parser):
        parser.add_argument(
            '--date',
            type=str,
            help='要分析的日期，格式为YYYY-MM-DD，默认为最近的交易日',
        )
        parser.add_argument(
            '--period',
            type=str,
            default='day',
            choices=['day', 'week', 'month'],
            help='分析周期，可选值：day, week, month，默认为day',
        )

    def handle(self, *args, **options):
        # 获取分析日期
        analysis_date = options.get('date')
        if analysis_date:
            analysis_date = datetime.datetime.strptime(analysis_date, '%Y-%m-%d').date()
        else:
            # 获取最近的交易日
            latest_daily = StockDaily.objects.order_by('-trade_date').first()
            if latest_daily:
                analysis_date = latest_daily.trade_date
            else:
                self.stdout.write(self.style.ERROR('没有找到交易数据'))
                return

        period = options.get('period', 'day')
        
        self.stdout.write(self.style.SUCCESS(f'开始分析 {analysis_date} 的市场数据，周期：{period}'))
        
        # 分析股票表现
        self.analyze_stock_performance(analysis_date, period)
        
        # 分析板块表现
        self.analyze_sector_performance(analysis_date, period)
        
        # 分析市场概览
        self.analyze_market_summary(analysis_date, period)
        
        self.stdout.write(self.style.SUCCESS('市场数据分析完成'))

    def analyze_stock_performance(self, analysis_date, period):
        """分析股票表现并保存结果"""
        self.stdout.write('分析股票表现...')
        
        # 清除当前日期和周期的旧数据
        StockPerformance.objects.filter(date=analysis_date, period=period).delete()
        
        # 获取当日交易数据
        daily_data = StockDaily.objects.filter(trade_date=analysis_date)
        
        # 获取股票基本信息
        stock_basics = {stock.ts_code: stock for stock in StockBasic.objects.all()}
        
        # 批量创建对象
        stock_performances = []
        
        for daily in daily_data:
            # 获取股票基本信息
            stock_basic = stock_basics.get(daily.ts_code)
            if not stock_basic:
                continue
                
            # 计算涨跌幅
            change_percent = ((daily.close - daily.open) / daily.open * 100) if daily.open > 0 else 0
            
            # 创建股票表现对象
            stock_performance = StockPerformance(
                stock_code=daily.ts_code,
                stock_name=stock_basic.name,
                date=analysis_date,
                period=period,
                open_price=daily.open,
                close_price=daily.close,
                high_price=daily.high,
                low_price=daily.low,
                volume=daily.vol,
                amount=daily.amount,
                change_percent=change_percent,
                turnover_rate=daily.turnover_rate if hasattr(daily, 'turnover_rate') else 0,
                sector=stock_basic.industry or '未知'
            )
            
            stock_performances.append(stock_performance)
            
        # 批量保存
        if stock_performances:
            StockPerformance.objects.bulk_create(stock_performances)
            self.stdout.write(self.style.SUCCESS(f'已保存 {len(stock_performances)} 条股票表现数据'))
        else:
            self.stdout.write(self.style.WARNING('没有找到股票表现数据'))

    def analyze_sector_performance(self, analysis_date, period):
        """分析板块表现并保存结果"""
        self.stdout.write('分析板块表现...')
        
        # 清除当前日期和周期的旧数据
        SectorPerformance.objects.filter(date=analysis_date, period=period).delete()
        
        # 获取当日所有股票表现数据
        stock_performances = StockPerformance.objects.filter(date=analysis_date, period=period)
        
        # 按板块分组统计
        sectors = {}
        
        for stock in stock_performances:
            sector = stock.sector
            if sector not in sectors:
                sectors[sector] = {
                    'stock_count': 0,
                    'up_count': 0,
                    'down_count': 0,
                    'changes': [],
                    'turnover_rates': []
                }
            
            sectors[sector]['stock_count'] += 1
            if stock.change_percent > 0:
                sectors[sector]['up_count'] += 1
            elif stock.change_percent < 0:
                sectors[sector]['down_count'] += 1
            
            sectors[sector]['changes'].append(stock.change_percent)
            sectors[sector]['turnover_rates'].append(stock.turnover_rate)
        
        # 创建板块表现对象
        sector_performances = []
        
        for sector_name, data in sectors.items():
            if not data['changes']:
                continue
                
            sector_performance = SectorPerformance(
                sector=sector_name,
                date=analysis_date,
                period=period,
                stock_count=data['stock_count'],
                up_count=data['up_count'],
                down_count=data['down_count'],
                avg_change=sum(data['changes']) / len(data['changes']),
                max_change=max(data['changes']),
                min_change=min(data['changes']),
                turnover_rate=sum(data['turnover_rates']) / len(data['turnover_rates']) if data['turnover_rates'] else 0
            )
            
            sector_performances.append(sector_performance)
        
        # 批量保存
        if sector_performances:
            SectorPerformance.objects.bulk_create(sector_performances)
            self.stdout.write(self.style.SUCCESS(f'已保存 {len(sector_performances)} 条板块表现数据'))
        else:
            self.stdout.write(self.style.WARNING('没有找到板块表现数据'))

    def analyze_market_summary(self, analysis_date, period):
        """分析市场概览并保存结果"""
        self.stdout.write('分析市场概览...')
        
        # 清除当前日期和周期的旧数据
        MarketSummary.objects.filter(date=analysis_date, period=period).delete()
        
        # 获取当日所有股票表现数据
        stock_performances = StockPerformance.objects.filter(date=analysis_date, period=period)
        
        # 统计市场数据
        total_stocks = stock_performances.count()
        up_stocks = stock_performances.filter(change_percent__gt=0).count()
        down_stocks = stock_performances.filter(change_percent__lt=0).count()
        flat_stocks = total_stocks - up_stocks - down_stocks
        
        # 计算平均涨跌幅和平均换手率
        avg_change = stock_performances.aggregate(avg=Avg('change_percent'))['avg'] or 0
        avg_turnover = stock_performances.aggregate(avg=Avg('turnover_rate'))['avg'] or 0
        
        # 计算总成交额
        total_amount = stock_performances.aggregate(sum=Sum('amount'))['sum'] or 0
        
        # 创建市场概览对象
        market_summary = MarketSummary(
            date=analysis_date,
            period=period,
            total_stocks=total_stocks,
            up_stocks=up_stocks,
            down_stocks=down_stocks,
            flat_stocks=flat_stocks,
            avg_change=avg_change,
            avg_turnover=avg_turnover,
            total_amount=total_amount
        )
        
        market_summary.save()
        self.stdout.write(self.style.SUCCESS('已保存市场概览数据'))